import pandas as pd
from sklearn.naive_bayes import MultinomialNB
from sklearn.feature_extraction.text import TfidfVectorizer
from py2neo import *


import jieba
import jieba.posseg as pseg
jieba.load_userdict("./data/userdict.txt")
import re
import sys,io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8")

import random

class QuestionPrediction():
    """
    问题预测: 采用NB进行多分类
    数据集：data/template_train.csv
    """
    def __init__(self):
        # 连接neo4j
        self.graph = Graph("http://localhost:7474",auth=("neo4j","neo4j"))

        # 编号和类别
        self.all_class = {
            0:'',
            1:'详细知识点',
            2:'后续',
            3:'',
            4:'',
            5:''
        }


        # 训练模板数据
        self.train_file = "./data/template_train.csv"

        # 词语重要度构建器
        self.tv = TfidfVectorizer()


        # 读取训练数据
        self.train_x, self.train_y = self.read_train_data(self.train_file)

        
        # 要训练的训练模型
        self.model = self.train_model_NB(self.train_x, self.train_y)


    # 从csv中抽取训练数据
    # csv有3列: label, text, desc
    # 其中desc没啥用, 只是用来说明的
    def read_train_data(self, template_train_file):
        train_data = pd.read_csv(template_train_file)
        train_x = train_data["text"].apply(lambda x: " ".join(list(jieba.cut(str(x))))).tolist()
        train_y = train_data["label"].tolist()
        return train_x,train_y


    # 采用 NB 训练模板数据
    def train_model_NB(self, X_train, y_train):
        train_data = self.tv.fit_transform(X_train).toarray()
        clf = MultinomialNB(alpha=0.01)
        clf.fit(train_data, y_train)
        return clf


    # 问题预测
    # 输入: 问题
    # 返回: 在neo4j中找到的答案
    def predict(self,question):
        res = pseg.cut(question)
        target = None


        word_container = []
        for word, flag in res:
            print("word:{}-flag:{}".format(word,flag))
            if(flag == 'nm'):
                target = word
                word_container.append(flag)
            else:
                word_container.append(word)

        if not target:
            return "没有识别到实体!!!"

        # print("关键词是:",target)
        question = [" ".join(word_container)]  #sk库设计的就是能处理好多条输入,但这里我们只需要一个
        print("question: {}".format(question[0]))
        test_data = self.tv.transform(question).toarray() #把这句话变成向量

        
        # 得到它是第几类问题
        class_num = self.model.predict(test_data)[0]
        class_name = self.all_class[class_num]
        
        print('问题的类别名字是: {}'.format(class_name))

        # 这里解决的是 某一tag的具体知识点  的问句
        if class_num == 1:
            # nodes = matcher.match("OJTag").where()
            cql = "match (tag:OJTag)-[r:Contains]->(op) where tag.name = '{}' return op".format(target)
            print('the cql is:',cql)

            str_container = []
            res = self.graph.run(cql).to_ndarray()
            # print(res.shape)
            # print(res.size)
            for i in range(res.size):
                node = res[i][0]

                name = node['name']
                description = node['description']
                description = re.sub("\|","\n\r",description)

                str_container.append(name)
                str_container.append(description)

            return "\n".join(str_container)

        if class_num == 2:
            return random.choice(['图','树','搜索','排序','哈希','动态规划'])


if __name__ == '__main__':
    question_model = QuestionPrediction()
    res = question_model.predict("数组学什么")
    print(res)

    
            


